# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http:#www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import typing



Expr = typing.Union[str, int, float, 'Expression']
PyInterval = typing.Union[str, 'IntervalUnit']
PyTimeUnit = typing.Union[str, 'TimeUnit']
#Use __ballista__ attribute to pass information about the object to the exhaustiveness script. Used for "skip" and "attr"
# "skip" tells the checker that these items should not be used when comparing to the items in the compiled extension

# "attr" combined with BallistaAttr marks the type declaration in a class as a class attribute for example
# class Expression:
#   expr: 'Expression' = 
#
#

#Errors being ignored below are about missing attributes on types
Expr.__ballista__ = "skip" #type: ignore
PyInterval.__ballista__ = "skip" #type: ignore
PyTimeUnit.__ballista__ = "skip"#type: ignore
typing.__ballista__ = "skip"#type: ignore



class TimeUnit: ...
class IntervalUnit: ...


class Expression:
    def __init__(self, expr: Expr)->None: ...
    def alias(self, name: str)->'Expression': ...
    def between(self, low: Expr, high: Expr, negated: bool=False)->'Expression':...
    def __add__(self, other: Expr) -> 'Expression': ...
    def __and__(self, other: Expr) -> 'Expression': ...
    #mypy warns when a non boolean type is returned from __ne__ or __eq__ operator 
    def __eq__(self, other: Expr) -> 'Expression': ... #type: ignore  
    def __ge__(self, other: Expr) -> 'Expression': ...
    def __gt__(self, other: Expr) -> 'Expression': ...
    def __invert__(self) -> 'Expression': ...
    def __le__(self, other: Expr) -> 'Expression': ...
    def __lt__(self, other: Expr) -> 'Expression': ...
    def __mul__(self, other: Expr) -> 'Expression': ...
    def __ne__(self, other: Expr) -> 'Expression': ... #type: ignore
    def __or__(self, other: Expr) -> 'Expression': ...
    def __radd__(self, other: Expr) -> 'Expression': ...
    def __rand__(self, other: Expr) -> 'Expression': ...
    def __rmul__(self, other: Expr) -> 'Expression': ...
    def __ror__(self, other: Expr) -> 'Expression': ...
    def __rsub__(self, other: Expr) -> 'Expression': ...
    def __rtruediv__(self, other: Expr) -> 'Expression': ...
    def __sub__(self, other: Expr) -> 'Expression': ...
    def __truediv__(self, other: Expr) -> 'Expression': ...

class Field:
    #While pyo3 defines the __new__ method mypy does not like that, use __init__ method instead
    def __init__(self, name:str, datatype: typing.Union[str, 'DataType'], nullable: bool = False, qualifier: typing.Optional[str] = None): ...
    def is_nullable(self)->bool: ...
    def datatype(self)->'DataType':...
    @staticmethod
    def from_json(json:str)->'Field':...
    def to_json(self)->str: ...

class DataType:
    def __init__(self, primitive: str):...
    def is_numeric(self)->bool: ...
    @staticmethod
    def from_json(json:str)->'DataType':...
    def to_json(self)->str: ...



Null = typing.Final[DataType]
Bool = typing.Final[DataType]
Int8 = typing.Final[DataType]
Int16 = typing.Final[DataType]
Int32 = typing.Final[DataType]
Int64 = typing.Final[DataType]
UInt8 = typing.Final[DataType]
UInt16 = typing.Final[DataType]
UInt32 = typing.Final[DataType]
UInt64 = typing.Final[DataType]
Float16 = typing.Final[DataType]
Float32 = typing.Final[DataType]
Float64 = typing.Final[DataType]
Binary = typing.Final[DataType]
LargeBinary = typing.Final[DataType]
Utf8 = typing.Final[DataType]
LargeUtf8 = typing.Final[DataType]



def List(item_type: Field)->DataType: ...
def LargeList(item_type: Field)->DataType: ...
def FixedSizeList(item_type: Field, list_size:int)->DataType:...
def FixedSizeBinary(binary_size:int)->DataType:...
def Interval(interval_unit: typing.Optional[PyInterval])->DataType: ...
def Date(date32: bool = False )->DataType: ...
def Time(time_unit: typing.Optional[PyTimeUnit], time32:bool = True)->DataType: ...
def Duration(time_unit: typing.Optional[PyTimeUnit])->DataType:...
def Timestamp(time_unit: typing.Optional[PyTimeUnit]=None, timezone: typing.Optional[str]=None)->DataType:...
def Union(*args: Field)->DataType: ...
def Struct(*args: Field)->DataType: ...
def Dictionary(key_type: DataType, value_type: DataType)->DataType: ...
def Decimal(whole: int, fractional: int)->DataType: ...

class CaseBuilder:
    #How to mark a class attribute in a way that mypy and the stub exhaustiveness checker both understand without polluting
    #the ballista namespace with marker classes
    expr: typing.Optional[Expression] = type("__attr_marker__", (object,), dict(__ballista__="attr")) #type: ignore
    def __init__(self, case_expr: typing.Optional[Expression] = None)->None:...
    def when(self, when_expr:Expr, then_expr: Expr)->'CaseBuilder':...
    def otherwise(self, otherwise_expr:Expr)->Expression: ...
    def build(self)->Expression: ...





def col(name:str)->Expression:...
def lit(literal: typing.Union[str, float, int])->Expression: ...
def sum(value: Expression)->Expression: ...
def min(value: Expression)->Expression: ...
def max(value: Expression)->Expression: ...
def count(value: Expression)->Expression: ...
def count_distinct(value: Expression)->Expression: ...
def concat(*args: Expression) -> Expression:...
def when(when: Expr, then: Expr)->CaseBuilder: ...
def avg(value: Expression)->Expression: ...
def case(case: Expr)->CaseBuilder: ...
#Unary macro generated functions
def array(*args: Expression)->Expression:...
def sqrt(unary: Expression)->Expression:...
def sin(unary: Expression)->Expression:...
def cos(unary: Expression)->Expression:...
def tan(unary: Expression)->Expression:...
def asin(unary: Expression)->Expression:...
def acos(unary: Expression)->Expression:...
def atan(unary: Expression)->Expression:...
def floor(unary: Expression)->Expression:...
def ceil(unary: Expression)->Expression:...
def round(unary: Expression)->Expression:...
def trunc(unary: Expression)->Expression:...
def abs(unary: Expression)->Expression:...
def signum(unary: Expression)->Expression:...
def exp(unary: Expression)->Expression:...
def ln(unary: Expression)->Expression:...
def log2(unary: Expression)->Expression:...
def log10(unary: Expression)->Expression:...
def lower(unary: Expression)->Expression:...
def trim(unary: Expression)->Expression:...
def upper(unary: Expression)->Expression:...


class Partitioning:

    @staticmethod
    def hash(n: int, *args: Expression)->'Partitioning':...

    @staticmethod
    def round_robin(n: int)->'Partitioning': ...


class Schema:
    def __init__(self, *args: Field, **kwargs: str): ...
    def to_json(self)->str: ...
    @staticmethod
    def from_json(json: str)->'Schema': ...

class DFSchema:
    def __init__(*args: typing.Union[Field, typing.Dict[str, typing.Any]], qualifier: typing.Optional[str], schema: typing.Optional[Schema]):...
    def join(self, schema: 'DFSchema')->'DFSchema': ...
class DataFrame:
    def select_columns(self, *args: str)->'DataFrame':...
    def select(self,  *args: Expr)->'DataFrame':...
    def filter(self,  *args: Expr)->'DataFrame':...
    def aggregate(self, group_expr: typing.List[Expr], aggr_expr: typing.List[Expr])->'DataFrame':...
    def limit(sefl, n: int)->'DataFrame':...
    def sort(self, *args:Expr)->'DataFrame':...
    def join(self, right: 'DataFrame', join_type:str, left_cols: typing.List[str], right_cols: typing.List[str])->'DataFrame':...
    def repartition(self, partitioning_scheme: Partitioning)->'DataFrame': ...
    #Figure out what to return here
    def collect(self)->typing.Any:...
    def explain(self, verbose: bool = False)->'DataFrame':...
    def schema(self)->DFSchema:...




class BallistaContext:
    def __init__(self, host: str = "localhost", port:int = 50051,  **kwds: str)->None:...
    def read_parquet(self, path:str)->None: ...
    def read_csv(self, name: str, path:str, has_header:bool = True, delimiter: str=",", file_extension: str = ".csv", schema: typing.Optional[Schema] = None, schema_infer_max_records: int = 1000 )->DataFrame:...
    def register_csv(self,name: str, path:str, has_header:bool = True, delimiter: str=",", file_extension: str = ".csv", schema: typing.Optional[Schema] = None, schema_infer_max_records: int = 1000)->None:...
    def register_table(self, name: str, table: DataFrame)->None: ...
    def register_parquet(self, name:str, path:str)->None: ...
    def sql(self, sql: str)->DataFrame: ...




